"""This module provides functionalities for structured logging, including:

- `logger`: a common logger
- `add_file_logger`: a function to add a file logging output to a logger
- `sm`: a helper class `StructuredMessage` to create structured logging style message
- `TimeIt`: a context manager to record running time for a piece of code
"""
import json
import logging
import time
from typing import Any

from common.config import SW_LOGLEVEL


LOGGER_NAME = "skywalker"
LOGGER_FORMAT = "%(asctime)s - %(name)s - %(levelname)s: %(message)s"
DEFAULT_LOGLEVEL = logging.INFO

KNOWN_LEVELS = {
    # unset, 0
    "debug": logging.DEBUG,  # 10
    "info": logging.INFO,  # 20
    "warn": logging.WARNING,  # 30
    "warning": logging.WARNING,  # 30
    # NOTE: warn and warning are the same, but after py3.7 warn is deprecated
    "error": logging.ERROR,  # 40
    # critical, 50
    # NOTE: we are not setting the `unset` and `critical` level
}

formatter = logging.Formatter(LOGGER_FORMAT)


def get_a_logger(name=None):
    level = KNOWN_LEVELS.get(SW_LOGLEVEL, DEFAULT_LOGLEVEL)
    logger = logging.getLogger(LOGGER_NAME if name is None else name)
    logger.setLevel(level)

    sh = logging.StreamHandler()
    sh.setFormatter(formatter)
    logger.addHandler(sh)
    return logger


# Optional: log to a file
def add_file_logger(logger: logging.Logger, fpath_log: str):
    """Log messages to a file along with logging to console.

    Parameters
    ----------
    logger: logging.Logger
        A logger object to be used.
    fpath_log: str
        File path to the logger file.

    Examples
    --------
    >>> from common.logger import logger, add_file_logger
    >>> logger = add_file_logger(logger, fpath_log='path/to/file.log')
    """
    fh = logging.FileHandler(fpath_log)
    fh.setFormatter(formatter)
    logger.addHandler(fh)
    return logger


class StructuredMessage(object):
    """A structured message with given message and more keyword args

    Parameters
    ----------
    message : str
        Message string.
    **kwargs : Dict[str, Any]
        Structured message organized as keyword arguments, assume the values can be serialized by json
    """

    def __init__(self, message: str, **kwargs: Any):
        self.message = message
        self.kwargs = kwargs

    def __str__(self):
        return "%s >>> %s" % (self.message, json.dumps(self.kwargs))


sm = StructuredMessage  # just a short name
logger = get_a_logger()  # create a default logger


class TimeIt(object):
    def __init__(self, description, how_to_report="print"):
        self.description = description
        mapping = {"print": 1, "ignore": 0, "logger": 2}
        if callable(how_to_report):
            self.how_to_report = how_to_report
        else:
            self.how_to_report = mapping[how_to_report]

    def __enter__(self):
        self.start = time.time()

    def __exit__(self, type, value, traceback):
        self.end = time.time()
        if self.how_to_report == 0:
            pass
        else:
            message = f"{self.description}: {self.end - self.start} sec"
            if self.how_to_report == 1:
                print(message)
            elif self.how_to_report == 2:
                logger.debug(message)
            else:
                self.how_to_report(message)


if __name__ == "__main__":
    # Use default logger only, log to console
    # sample usage of the logger and sm (StructuredMessage)
    logger.debug(
        sm("message 1 some debug info", foo="bar", bar="baz", num=123, fnum=123.456)
    )
    logger.info("message 2 info we are doing something")
    logger.warning(
        sm("message 3 warning we might get something wrong!", critical_property=10000)
    )
    logger.error("message 4 mayday and we are quitting")

    # Update logger to log to both the console and a file.
    logger1 = get_a_logger("skywalker-showcase")
    logger1 = add_file_logger(logger1, fpath_log="./file.log")
    logger1.info("message log to the console and a file.")

    # use the TimeIt contextmanager
    with TimeIt("just a sleepy test", how_to_report="print"):
        time.sleep(1)

    with TimeIt("another sleepy test using Skywalker Logger", how_to_report="logger"):
        time.sleep(0.5)

    with TimeIt(
        "sleepy test using a logging function",
        how_to_report=logger.info,
    ):
        time.sleep(0.5)
